{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 0
   },
   "source": [
    "# Matrix Multiplication\n",
    "\n",
    "Matrix Multiplication is one of the most widely operators in scientific computing and deep learning, which is typically referred to as *GEMM* (GEneral Matrix Multiply). Let's implement its computation in this section.\n",
    "\n",
    "Given $A\\in\\mathbb R^{n\\times l}$, and $B \\in\\mathbb R^{l\\times m}$, if $C=AB$ then $C \\in\\mathbb R^{n\\times m}$ and\n",
    "\n",
    "$$C_{i,j} = \\sum_{k=1}^l A_{i,k} B_{k,j}.$$\n",
    "\n",
    "The elements accessed to compute $C_{i,j}$ are illustrated in :numref:`fig_matmul_default`.\n",
    "\n",
    "![Compute $C_{x,y}$ in matrix multiplication.](../img/matmul_default.svg)\n",
    ":label:`fig_matmul_default`\n",
    "\n",
    "The following method returns the computing expression of matrix multiplication.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "1"
    },
    "origin_pos": 1,
    "tab": [
     "tvm"
    ]
   },
   "outputs": [],
   "source": [
    "import d2ltvm\n",
    "import numpy as np\n",
    "import tvm\n",
    "from tvm import te\n",
    "\n",
    "# Save to the d2ltvm package\n",
    "def matmul(n, m, l):\n",
    "    \"\"\"Return the computing expression of matrix multiplication\n",
    "    A : n x l matrix\n",
    "    B : l x m matrix\n",
    "    C : n x m matrix with C = A B\n",
    "    \"\"\"\n",
    "    k = te.reduce_axis((0, l), name='k')\n",
    "    A = te.placeholder((n, l), name='A')\n",
    "    B = te.placeholder((l, m), name='B')\n",
    "    C = te.compute((n, m),\n",
    "                    lambda x, y: te.sum(A[x, k] * B[k, y], axis=k),\n",
    "                    name='C')\n",
    "    return A, B, C"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 2
   },
   "source": [
    "Let's compile a module for a square matrix multiplication.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "2"
    },
    "origin_pos": 3,
    "tab": [
     "tvm"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "// attr [C] storage_scope = \"global\"\n",
      "allocate C[float32 * 10000]\n",
      "produce C {\n",
      "  for (x, 0, 100) {\n",
      "    for (y, 0, 100) {\n",
      "      C[((x*100) + y)] = 0f\n",
      "      for (k, 0, 100) {\n",
      "        C[((x*100) + y)] = (C[((x*100) + y)] + (A[((x*100) + k)]*B[((k*100) + y)]))\n",
      "      }\n",
      "    }\n",
      "  }\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "n = 100\n",
    "A, B, C = matmul(n, n, n)\n",
    "s = te.create_schedule(C.op)\n",
    "print(tvm.lower(s, [A, B], simple_mode=True))\n",
    "mod = tvm.build(s, [A, B, C])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 4
   },
   "source": [
    "The pseudo code is simply a naive 3-level nested for loop to calculate the matrix multiplication.\n",
    "\n",
    "And then we verify the results. Note that NumPy may use multi-threading to accelerate its computing, which may result in slightly different results due to the numerical error. There we use `assert_allclose` with a relative large tolerant error to test the correctness.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "3"
    },
    "origin_pos": 5,
    "tab": [
     "tvm"
    ]
   },
   "outputs": [],
   "source": [
    "a, b, c = d2ltvm.get_abc((100, 100), tvm.nd.array)\n",
    "mod(a, b, c)\n",
    "np.testing.assert_allclose(np.dot(a.asnumpy(), b.asnumpy()),\n",
    "                           c.asnumpy(), atol=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 6
   },
   "source": [
    "## Summary\n",
    "\n",
    "- We can express the computation of matrix multiplication in TVM in one line of code.\n",
    "- The naive matrix multiplication is a 3-level nested for loop.\n"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}